import os

from PIL import Image
from loguru import logger


async def compress_image(input_path: str, output_path: str, target_size: int):
    with Image.open(input_path) as img:
        # 如果图像是 RGBA 模式，则转换为 RGB 模式
        if img.mode in ('RGBA', 'LA') or (img.mode == 'P' and 'transparency' in img.info):
            img = img.convert('RGB')

        # 灰度化
        img = img.convert('L')

        for i in range(1, 19):
            quality = (10 - i) * 10 if i < 10 else 19 - i
            img.save(output_path, "JPEG", quality=quality)
            pic_size = os.path.getsize(output_path)
            if pic_size < target_size:
                break

    logger.info(f'compress_image({input_path} -> {output_path}) success')


async def merge_images_horizontally(img_paths: list[str], output_path: str):
    images = [Image.open(x) for x in img_paths]

    # 计算总宽度和最大高度
    widths, heights = zip(*(i.size for i in images))
    total_width = sum(widths)
    max_height = max(heights)

    new_img = Image.new('RGB', (total_width, max_height))

    x_offset = 0
    for img in images:
        new_img.paste(img, (x_offset, 0))
        x_offset += img.width

    new_img.save(output_path)


async def merge_images_vertically(img_paths: list[str], output_path: str):
    images = [Image.open(x) for x in img_paths]

    # 计算总高度和最大宽度
    widths, heights = zip(*(i.size for i in images))
    total_height = sum(heights)
    max_width = max(widths)

    new_img = Image.new('RGB', (max_width, total_height))

    y_offset = 0
    for img in images:
        new_img.paste(img, (0, y_offset))
        y_offset += img.height

    new_img.save(output_path)
